#include "TCPFilter.h"
#include <iostream>
#include <cstring>
#include <algorithm>
#include <csignal>
#include <arpa/inet.h>
#include <thread>
#include <chrono>

// 构造函数
TCPPacketFilter::TCPPacketFilter(const std::string& dev) 
    : handle(nullptr), running(false), device(dev),
      syn_threshold(20), time_window(1), block_duration(60) {}

// 析构函数
TCPPacketFilter::~TCPPacketFilter() {
    stopCapture();
}

// 添加过滤规则
void TCPPacketFilter::addRule(uint32_t src_ip, uint32_t dst_ip, 
            uint16_t src_port, uint16_t dst_port,
            uint8_t flags, bool drop, bool log) {
    rules.push_back({src_ip, dst_ip, src_port, dst_port, flags, drop, log});
}

// 初始化libpcap
bool TCPPacketFilter::init() {
    char errbuf[PCAP_ERRBUF_SIZE];
    
    // 打开网络设备
    handle = pcap_open_live(device.c_str(), BUFSIZ, 1, 1000, errbuf);
    if (handle == nullptr) {
        std::cerr << "打开设备失败: " << errbuf << std::endl;
        return false;
    }
    
    // 设置过滤表达式，只捕获TCP包
    struct bpf_program fp;
    const char* filter_exp = "tcp";
    if (pcap_compile(handle, &fp, filter_exp, 0, PCAP_NETMASK_UNKNOWN) == -1) {
        std::cerr << "编译过滤器失败: " << pcap_geterr(handle) << std::endl;
        return false;
    }
    
    if (pcap_setfilter(handle, &fp) == -1) {
        std::cerr << "设置过滤器失败: " << pcap_geterr(handle) << std::endl;
        pcap_freecode(&fp);
        return false;
    }
    
    pcap_freecode(&fp);
    return true;
}

// 检查数据包是否匹配规则
bool TCPPacketFilter::matchRule(const iphdr* ip_hdr, const tcphdr* tcp_hdr) const {
    for (const auto& rule : rules) {
        // 检查IP地址
        if (rule.src_ip != 0 && rule.src_ip != ip_hdr->saddr) continue;
        if (rule.dst_ip != 0 && rule.dst_ip != ip_hdr->daddr) continue;
        
        // 检查端口
        if (rule.src_port != 0 && rule.src_port != ntohs(tcp_hdr->source)) continue;
        if (rule.dst_port != 0 && rule.dst_port != ntohs(tcp_hdr->dest)) continue;
        
        // 检查TCP标志位
        if (rule.flags != 0 && (tcp_hdr->th_flags & rule.flags) != rule.flags) continue;
        
        return true;
    }
    return false;
}

// 获取匹配规则的动作
FilterRule TCPPacketFilter::getMatchedRule(const iphdr* ip_hdr, const tcphdr* tcp_hdr) const {
    for (const auto& rule : rules) {
        if (rule.src_ip != 0 && rule.src_ip != ip_hdr->saddr) continue;
        if (rule.dst_ip != 0 && rule.dst_ip != ip_hdr->daddr) continue;
        if (rule.src_port != 0 && rule.src_port != ntohs(tcp_hdr->source)) continue;
        if (rule.dst_port != 0 && rule.dst_port != ntohs(tcp_hdr->dest)) continue;
        if (rule.flags != 0 && (tcp_hdr->th_flags & rule.flags) != rule.flags) continue;
        
        return rule;
    }
    return FilterRule{0, 0, 0, 0, 0, false, false}; // 默认规则
}

// SYN洪水防御处理
bool TCPPacketFilter::handleSYNFlood(uint32_t src_ip) {
    std::lock_guard<std::mutex> lock(flood_mutex);
    time_t now = time(nullptr);
    
    // 检查IP是否已被封锁
    auto it = syn_flood_map.find(src_ip);
    if (it != syn_flood_map.end() && it->second.blocked) {
        // 检查是否超过封锁时间
        if (now - it->second.block_time > block_duration) {
            // 解封IP
            char ip_str[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &src_ip, ip_str, INET_ADDRSTRLEN);
            std::cout << "自动解封IP: " << ip_str << std::endl;
            syn_flood_map.erase(it);
            return false;
        }
        return true; // 继续封锁
    }
    
    // 初始化或更新状态
    if (it == syn_flood_map.end()) {
        syn_flood_map[src_ip] = {now, 1, false, 0};
    } else {
        // 检查时间窗口
        if (now - it->second.last_syn_time > time_window) {
            // 超过时间窗口，重置计数
            it->second.last_syn_time = now;
            it->second.syn_count = 1;
        } else {
            // 在时间窗口内，增加计数
            it->second.syn_count++;
            it->second.last_syn_time = now;
            
            // 检查是否超过阈值
            if (it->second.syn_count >= syn_threshold) {
                // 触发封锁
                it->second.blocked = true;
                it->second.block_time = now;
                
                char ip_str[INET_ADDRSTRLEN];
                inet_ntop(AF_INET, &src_ip, ip_str, INET_ADDRSTRLEN);
                std::cout << "检测到SYN洪水攻击! 封锁IP: " << ip_str 
                          << " (" << it->second.syn_count 
                          << " SYN包/" << time_window << "秒)" << std::endl;
                
                // 添加动态规则封锁此IP
                addRule(src_ip, 0, 0, 0, TH_SYN, true, true);
                
                return true;
            }
        }
    }
    
    return false;
}

// 包处理回调函数
void TCPPacketFilter::packetHandler(u_char *userData, const struct pcap_pkthdr *pkthdr, const u_char *packet) {
    TCPPacketFilter* filter = reinterpret_cast<TCPPacketFilter*>(userData);
    
    // 跳过以太网头（14字节）
    const struct iphdr* ip_hdr = reinterpret_cast<const struct iphdr*>(packet + 14);
    unsigned int ip_header_len = ip_hdr->ihl * 4;
    
    // 确保数据包包含完整的TCP头部
    if (pkthdr->len < 14 + ip_header_len + sizeof(tcphdr)) {
        return;
    }
    
    // 解析TCP头
    const struct tcphdr* tcp_hdr = reinterpret_cast<const struct tcphdr*>(packet + 14 + ip_header_len);
    
    // 获取源和目的IP地址
    char src_ip_str[INET_ADDRSTRLEN];
    char dst_ip_str[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(ip_hdr->saddr), src_ip_str, INET_ADDRSTRLEN);
    inet_ntop(AF_INET, &(ip_hdr->daddr), dst_ip_str, INET_ADDRSTRLEN);
    
    // SYN洪水防御检查 (仅处理SYN包)
    if (tcp_hdr->th_flags & TH_SYN && !(tcp_hdr->th_flags & TH_ACK)) {
        if (filter->handleSYNFlood(ip_hdr->saddr)) {
            // 此IP已被封锁，记录并返回
            std::cout << "拦截SYN洪水攻击包: " 
                      << src_ip_str << ":" << ntohs(tcp_hdr->source)
                      << " -> " << dst_ip_str << ":" << ntohs(tcp_hdr->dest)
                      << std::endl;
            return;
        }
    }
    
    // 检查是否匹配规则
    FilterRule matched_rule = filter->getMatchedRule(ip_hdr, tcp_hdr);
    
    // 记录日志
    if (matched_rule.log) {
        std::cout << "\nTCP包: "
                  << src_ip_str << ":" << ntohs(tcp_hdr->source)
                  << " -> " << dst_ip_str << ":" << ntohs(tcp_hdr->dest)
                  << " 标志: 0x" << std::hex << static_cast<int>(tcp_hdr->th_flags) << std::dec
                  << " 大小: " << pkthdr->len << " 字节"
                  << " 动作: " << (matched_rule.drop ? "丢弃" : "允许");
    }
    
    if (matched_rule.drop) {
        // 丢弃数据包
    }
}

// 开始捕获数据包
void TCPPacketFilter::startCapture() {
    if (handle == nullptr) {
        std::cerr << "未初始化libpcap" << std::endl;
        return;
    }
    
    running = true;
    std::cout << "开始捕获TCP数据包 (按Ctrl+C停止)..." << std::endl;
    
    // 捕获数据包
    pcap_loop(handle, 0, packetHandler, reinterpret_cast<u_char*>(this));
}

// 停止捕获
void TCPPacketFilter::stopCapture() {
    running = false;
    if (handle) {
        pcap_breakloop(handle); // 中断捕获循环
        pcap_close(handle);
        handle = nullptr;
    }
}

// 获取可用网络设备
void TCPPacketFilter::listDevices() {
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_if_t *alldevs;
    
    if (pcap_findalldevs(&alldevs, errbuf) == -1) {
        std::cerr << "查找设备失败: " << errbuf << std::endl;
        return;
    }
    
    std::cout << "\n可用网络设备:" << std::endl;
    for (pcap_if_t *d = alldevs; d != nullptr; d = d->next) {
        std::cout << "  - " << d->name;
        if (d->description)
            std::cout << " (" << d->description << ")";
        std::cout << std::endl;
    }
    
    pcap_freealldevs(alldevs);
}

// 获取当前被封锁的IP列表
std::vector<uint32_t> TCPPacketFilter::getBlockedIPs() const {
    std::vector<uint32_t> blocked_ips;
    std::lock_guard<std::mutex> lock(flood_mutex);
    
    for (const auto& entry : syn_flood_map) {
        if (entry.second.blocked) {
            blocked_ips.push_back(entry.first);
        }
    }
    
    return blocked_ips;
}

// 手动封锁IP
void TCPPacketFilter::blockIP(uint32_t ip) {
    std::lock_guard<std::mutex> lock(flood_mutex);
    
    // 如果IP不在映射中，添加它
    if (syn_flood_map.find(ip) == syn_flood_map.end()) {
        syn_flood_map[ip] = {time(nullptr), 0, true, time(nullptr)};
    } else {
        syn_flood_map[ip].blocked = true;
        syn_flood_map[ip].block_time = time(nullptr);
    }
    
    // 添加规则封锁此IP的所有SYN包
    addRule(ip, 0, 0, 0, TH_SYN, true, true);
    
    char ip_str[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &ip, ip_str, INET_ADDRSTRLEN);
    std::cout << "手动封锁IP: " << ip_str << std::endl;
}

// 手动解封IP
void TCPPacketFilter::unblockIP(uint32_t ip) {
    std::lock_guard<std::mutex> lock(flood_mutex);
    
    auto it = syn_flood_map.find(ip);
    if (it != syn_flood_map.end()) {
        it->second.blocked = false;
        
        // 移除封锁规则
        for (auto rule_it = rules.begin(); rule_it != rules.end(); ) {
            if (rule_it->src_ip == ip && rule_it->flags == TH_SYN && rule_it->drop) {
                rule_it = rules.erase(rule_it);
            } else {
                ++rule_it;
            }
        }
        
        char ip_str[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &ip, ip_str, INET_ADDRSTRLEN);
        std::cout << "手动解封IP: " << ip_str << std::endl;
    }
}
// ... 现有 TCPFilter.cpp 代码保持不变 ...

// 新增的封装函数
void runTCPFilter(const std::string& device) {
    // 创建过滤器
    TCPPacketFilter filter(device);
    
    // 初始化libpcap
    if (!filter.init()) {
        std::cerr << "初始化失败" << std::endl;
        return;
    }
    
    // 设置SYN洪水防御参数
    filter.setSYNFloodParams(20, 1, 60);
    
    // 添加过滤规则
    filter.addRule(0, inet_addr("192.168.1.100"), 0, 80, TH_SYN, true); // 丢弃所有到192.168.1.100:80的SYN包
    filter.addRule(inet_addr("10.0.0.5"), 0, 0, 22, 0, false);         // 记录来自10.0.0.5到任意主机22端口的包
    filter.addRule(0, 0, 0, 443, 0, false);                           // 记录所有HTTPS流量
    
    std::cout << "TCP包过滤器已启动，使用网卡: " << device << std::endl;
    std::cout << "使用以下规则：" << std::endl;
    std::cout << "1. 丢弃所有到 192.168.1.100:80 的SYN包" << std::endl;
    std::cout << "2. 记录来自 10.0.0.5 到任意主机22端口的包" << std::endl;
    std::cout << "3. 记录所有HTTPS流量 (端口443)" << std::endl;
    std::cout << "SYN洪水防御已启用: 阈值=20 SYN包/秒, 封锁时间=60秒" << std::endl;
    
    // 开始捕获
    filter.startCapture();
}
